﻿#include <iostream>
#include <vector>
#include <cmath>
#include <iomanip>
#include <fstream>
#include <nlohmann/json.hpp>
#include "../include/vector.h"
#include "../include/convolution.h"
#include "../include/iir.h"

using namespace std;
using namespace splab;
using json = nlohmann::json;

// 实现零相位滤波函数
template<typename Type>
Vector<Type> filtfilt(const Vector<Type>& signal, const Vector<Type>& numCoefs, const Vector<Type>& denCoefs) {
    // 正向滤波
    Vector<Type> forward = convolution(signal, numCoefs);
    forward = convolution(forward, denCoefs);
    
    // 反转信号
    Vector<Type> reversed = reverse(forward);
    
    // 反向滤波
    Vector<Type> backward = convolution(reversed, numCoefs);
    backward = convolution(backward, denCoefs);
    
    // 再次反转得到最终结果
    Vector<Type> result = reverse(backward);
    
    // 如果结果长度大于输入信号长度，只取前signal.size()个点
    if (result.size() > signal.size()) {
        Vector<Type> truncated(signal.size());
        for (int i = 0; i < signal.size(); ++i) {
            truncated[i] = result[i];
        }
        return truncated;
    }
    
    return result;
}

// 生成测试信号：包含多个频率成分
template<typename Type>
Vector<Type> generateTestSignal(int length, Type fs, const vector<Type>& freqs, const vector<Type>& amplitudes) {
    Vector<Type> signal(length);
    Type t;
    
    for(int i = 0; i < length; i++) {
        t = Type(i) / fs;
        signal[i] = 0;  // 初始化为0
        
        // 生成多个频率的正弦波叠加
        for(size_t j = 0; j < freqs.size(); j++) {
            signal[i] += amplitudes[j] * sin(2 * PI * freqs[j] * t);
        }
    }
    
    return signal;
}

int main() {
    // 设置采样频率和信号长度
    double fs = 250.0;  // 采样频率250Hz
    int signalLength = 1000;  // 信号长度1000点
    
    // 设置频率和幅度
    vector<double> freqs = {5.0, 15.0, 25.0};  // 频率列表：5Hz, 15Hz, 25Hz
    vector<double> amplitudes = {1.0, 1.3, 1.5};  // 对应频率的幅度
    
    // 生成测试信号
    Vector<double> signal = generateTestSignal<double>(signalLength, fs, freqs, amplitudes);
    
    // 设计低通IIR滤波器（Butterworth）
    IIR lowpassFilter("lowpass", "Butterworth");
    // 设置参数：采样率，通带频率，通带波纹，阻带频率，阻带衰减
    lowpassFilter.setParams(fs, 20.0, -3.0, 21.0, -40.0);  // 截止频率20Hz
    lowpassFilter.design();
    
    // 设计高通IIR滤波器（Butterworth）
    IIR highpassFilter("highpass", "Butterworth");
    // 设置参数：采样率，阻带频率，阻带衰减，通带频率，通带波纹
    highpassFilter.setParams(fs, 9.0, -40.0, 10.0, -3.0);  // 截止频率10Hz
    highpassFilter.design();
    
    // 获取滤波器系数
    Vector<double> lowpassNumCoefs = lowpassFilter.getNumCoefs();
    Vector<double> lowpassDenCoefs = lowpassFilter.getDenCoefs();
    Vector<double> highpassNumCoefs = highpassFilter.getNumCoefs();
    Vector<double> highpassDenCoefs = highpassFilter.getDenCoefs();
    
    // 输出滤波器系数
    cout << "\n低通滤波器系数：" << endl;
    cout << "分子系数 (b):" << endl;
    for(int i = 0; i < lowpassNumCoefs.size(); i++) {
        cout << "b[" << i << "] = " << scientific << setprecision(6) << lowpassNumCoefs[i] << endl;
    }
    cout << "\n分母系数 (a):" << endl;
    for(int i = 0; i < lowpassDenCoefs.size(); i++) {
        cout << "a[" << i << "] = " << scientific << setprecision(6) << lowpassDenCoefs[i] << endl;
    }
    
    cout << "\n高通滤波器系数：" << endl;
    cout << "分子系数 (b):" << endl;
    for(int i = 0; i < highpassNumCoefs.size(); i++) {
        cout << "b[" << i << "] = " << scientific << setprecision(6) << highpassNumCoefs[i] << endl;
    }
    cout << "\n分母系数 (a):" << endl;
    for(int i = 0; i < highpassDenCoefs.size(); i++) {
        cout << "a[" << i << "] = " << scientific << setprecision(6) << highpassDenCoefs[i] << endl;
    }
    
    // 先应用低通滤波
    Vector<double> filteredSignal = signal;
    // filteredSignal = filtfilt(filteredSignal, lowpassNumCoefs, lowpassDenCoefs);
    // 再应用高通滤波
    filteredSignal = filtfilt(filteredSignal, highpassNumCoefs, highpassDenCoefs);
    
    // 保存原始信号到JSON文件
    json originData;
    originData["sampling_rate"] = fs;
    originData["data_length"] = signalLength;
    originData["data"] = vector<double>(signal.begin(), signal.end());
    
    ofstream originFile("origin_data.json");
    originFile << std::setw(4) << originData << endl;
    originFile.close();
    
    // 保存滤波后信号到JSON文件
    json filteredData;
    filteredData["sampling_rate"] = fs;
    filteredData["data_length"] = signalLength;
    filteredData["data"] = vector<double>(filteredSignal.begin(), filteredSignal.end());
    
    ofstream filteredFile("filtered_data.json");
    filteredFile << std::setw(4) << filteredData << endl;
    filteredFile.close();
    
    // 输出结果
    cout << "原始信号和滤波后信号的前20个采样点：" << endl;
    cout << setw(10) << "时间(s)" << setw(15) << "原始信号" << setw(15) << "滤波后信号" << endl;
    cout << "------------------------------------------------" << endl;
    
    for(int i = 0; i < 20; i++) {
        cout << setw(10) << fixed << setprecision(3) << i/fs
             << setw(15) << signal[i]
             << setw(15) << filteredSignal[i] << endl;
    }
    
    cout << "\n数据已保存到 origin_data.json 和 filtered_data.json" << endl;
    
    return 0;
} 